import mediapipe as mp
import numpy as np
import cv2

mp_pose = mp.solutions.pose
pose = mp_pose.Pose(model_complexity=1)

def get_landmark_coords(results):
    """提取33个关键点的坐标，返回字典 { landmark_id: (x,y,z) }"""
    coords = {}
    if results.pose_landmarks:
        for idx, landmark in enumerate(results.pose_landmarks.landmark):
            coords[idx] = (landmark.x, landmark.y, landmark.z)  # x,y,z为归一化坐标
    return coords

# 原estimate_pose函数修改为返回帧、关键点坐标、原始results
def estimate_pose(frame):
    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    results = pose.process(frame_rgb)
    # 绘制关键点
    frame = cv2.cvtColor(frame_rgb, cv2.COLOR_RGB2BGR)
    if results.pose_landmarks:
        mp.solutions.drawing_utils.draw_landmarks(
            frame, results.pose_landmarks, mp_pose.POSE_CONNECTIONS
        )
    # 提取坐标
    coords = get_landmark_coords(results)
    return frame, coords, results
from scipy.signal import savgol_filter  # 需要在requirements.txt加scipy

class SGFilter:
    def __init__(self, window_size=5, polyorder=2):
        self.window_size = window_size  # 滑动窗口大小（奇数）
        self.polyorder = polyorder  # 多项式阶数
        self.history = {}  # 缓存历史关键点数据 { landmark_id: [ (x1,y1,z1), (x2,y2,z2), ... ] }

    def smooth(self, coords):
        """对当前帧的关键点坐标进行平滑，返回平滑后的数据"""
        smoothed_coords = {}
        for idx, (x, y, z) in coords.items():
            # 缓存历史数据（保留window_size-1帧）
            if idx not in self.history:
                self.history[idx] = []
            self.history[idx].append([x, y, z])
            self.history[idx] = self.history[idx][-(self.window_size-1):]  # 截断
            # 拼接当前帧数据，形成窗口
            window = self.history[idx] + [[x, y, z]]
            if len(window) >= self.window_size:
                # 对x,y,z分别滤波
                window_np = np.array(window)
                smoothed = savgol_filter(window_np, self.window_size, self.polyorder, axis=0)
                smoothed_coords[idx] = tuple(smoothed[-1])  # 取最后一帧（当前帧）的平滑结果
            else:
                smoothed_coords[idx] = (x, y, z)  # 窗口不足时返回原始值
        return smoothed_coords